from concurrent import futures
import grpc
import rag_agent_pb2
import rag_agent_pb2_grpc

from ai_service import run_rag_agent


class RagServiceServicer(rag_agent_pb2_grpc.RagServiceServicer):
    def AskWithContent(self, request, context):
        for msg in run_rag_agent(request.query, request.content):
            yield rag_agent_pb2.Answer(answer=msg)

def serve():
    server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
    rag_agent_pb2_grpc.add_RagServiceServicer_to_server(
        RagServiceServicer(), server
    )
    server.add_insecure_port("[::]:50051")
    server.start()
    server.wait_for_termination()

if __name__ == "__main__":
    print("Hello from ai-service!")
    serve()
